{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "X len: 1180 Y len: 1180\n",
      "[[9, 15, 36, 15, 36, 34, 31, 27, 27, 26, 29, 37, 23, 25, 34, 15, 31, 36, 33, 34, 17, 38, 23, 23, 11, 37, 20, 34, 1, 15, 36, 38, 23, 33, 17, 7, 25, 8, 17, 15, 17, 23, 25, 27, 36, 33, 25, 6, 9, 42, 23, 26, 37, 25, 1, 13, 1, 38, 21, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]\n",
      "[2]\n",
      "Dectected cnt: 1180 total: 1180\n",
      "Dectect Rate is: 1.0\n"
     ]
    }
   ],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "#from list2array2matrix2input2softmax import *\n",
    "import os\n",
    "import pickle\n",
    "from keras.utils import to_categorical\n",
    "from keras.models import load_model\n",
    "import numpy as np\n",
    "SRC_IP_IDX = 3-1               \n",
    "DST_IP_IDX = 4-1               \n",
    "SRC_PORT_IDX = 5-1             \n",
    "DST_PORT_IDX = 6-1             \n",
    "PROTOCOL_IDX = 7-1             \n",
    "DNS_QUERY_NAME_IDX = 55-1 # domain\n",
    "DNS_REQUEST_TYPE = 56-1\n",
    "DNS_DOMAIN_TTL = 59-1\n",
    "DNS_REPLY_IPV4IP = 60-1        \n",
    "DNS_REPLY_IPV6IP = 61-1        \n",
    "DNS_REPLY_RRTYPE = 62-1        \n",
    "DNS_REQUEST_LEN  = 88-1        \n",
    "DNS_REPLY_LENGTH = 90-1  \n",
    "\n",
    "def pad_sequences(X, maxlen, value=0):\n",
    "    S=[]\n",
    "    for x in X:\n",
    "        xlen = len(x)\n",
    "        if xlen < maxlen:\n",
    "            x.extend([value]*(maxlen-xlen))\n",
    "        else:\n",
    "            x = x[:maxlen]\n",
    "        S.append(x)\n",
    "    return S\n",
    "def extract_domain(domain):\n",
    "    suffix = {'.com','.la','.io', '.co', '.cn','.info', '.net', '.org','.me', '.mobi', '.us', '.biz', '.xxx', '.ca', '.co.jp', '.com.cn', '.net.cn', '.org.cn', '.mx','.tv', '.ws', '.ag', '.com.ag', '.net.ag', '.org.ag','.am','.asia', '.at', '.be', '.com.br', '.net.br', '.name', '.live', '.news', '.bz', '.tech', '.pub', '.wang', '.space', '.top', '.xin', '.social', '.date', '.site', '.red', '.studio', '.link', '.online', '.help', '.kr', '.club', '.com.bz', '.net.bz', '.cc', '.band', '.market', '.com.co', '.net.co', '.nom.co', '.lawyer', '.de', '.es', '.com.es', '.nom.es', '.org.es', '.eu', '.wiki', '.design', '.software', '.fm', '.fr', '.gs', '.in', '.co.in', '.firm.in', '.gen.in', '.ind.in', '.net.in', '.org.in', '.it', '.jobs', '.jp', '.ms', '.com.mx', '.nl','.nu','.co.nz','.net.nz', '.org.nz', '.se', '.tc', '.tk', '.tw', '.com.tw', '.idv.tw', '.org.tw', '.hk', '.co.uk', '.me.uk', '.org.uk', '.vg'}\n",
    "\n",
    "    domain = domain.lower()\n",
    "    names = domain.split(\".\")\n",
    "    if len(names) >= 3:\n",
    "        if (\".\"+\".\".join(names[-2:])) in suffix:\n",
    "            return \".\".join(names[-3:]), \".\".join(names[:-3])\n",
    "        elif (\".\"+names[-1]) in suffix:\n",
    "            return \".\".join(names[-2:]), \".\".join(names[:-2])\n",
    "    #print (\"New domain suffix found. Use tld extract domain...\")\n",
    "\n",
    "    pos = domain.rfind(\"/\")\n",
    "    if pos >= 0: # maybe subdomain contains /, for dns tunnel tool\n",
    "        ext = tldextract.extract(domain[pos+1:])\n",
    "        subdomain = domain[:pos+1] + ext.subdomain\n",
    "    else:\n",
    "        ext = tldextract.extract(domain)\n",
    "        subdomain = ext.subdomain\n",
    "    if ext.suffix:\n",
    "        mdomain = ext.domain + \".\" + ext.suffix\n",
    "    else:\n",
    "        mdomain = ext.domain\n",
    "    return mdomain, subdomain\n",
    "\n",
    "\n",
    "def iterbrowse(path):\n",
    "    for home, dirs, files in os.walk(path):\n",
    "        for filename in files:\n",
    "            yield os.path.join(home, filename)\n",
    "            \n",
    "def metadata2_domain_data(log): \n",
    "    data = log.split('^')\n",
    "    if not filter_metadata_dns(data):\n",
    "        return None, None\n",
    "    domain = data[DNS_QUERY_NAME_IDX]\n",
    "    mdomain, subdomain = extract_domain(domain)\n",
    "    return (mdomain, subdomain)\n",
    "\n",
    "def filter_metadata_dns(data):\n",
    "    if(len(data) < 91):\n",
    "        return False\n",
    "\n",
    "    protol  = data[PROTOCOL_IDX]\n",
    "    dstport = data[DST_PORT_IDX]\n",
    "    dstip   = data[DST_IP_IDX]\n",
    "    qname   = data[DNS_QUERY_NAME_IDX]\n",
    "\n",
    "    if '' == qname or '' == dstip:\n",
    "        return False\n",
    "    if '17' == protol and ('53' == dstport):\n",
    "        return True\n",
    "    return False\n",
    "\n",
    "class LABEL(object):\n",
    "    white = 0\n",
    "    cdn = 1\n",
    "    black = 2\n",
    "def get_predict_data():\n",
    "    data_path = \"./xshell_data\"\n",
    "    black_data = []\n",
    "    for path in iterbrowse(data_path):\n",
    "        with open(path) as f:\n",
    "            for line in f:\n",
    "                mdomain, subdomain = metadata2_domain_data(line)\n",
    "                if subdomain is not None:\n",
    "                    black_data.append(subdomain)\n",
    "    return black_data\n",
    "\n",
    "\n",
    "org_X = []\n",
    "\n",
    "def get_xshell_data():\n",
    "    global org_X\n",
    "    org_X = get_predict_data()\n",
    "    labels = [LABEL.black]*len(org_X)\n",
    "\n",
    "    volcab_file = \"volcab.pkl\"\n",
    "    assert os.path.exists(volcab_file)\n",
    "    pkl_file = open(volcab_file, 'rb')\n",
    "    data = pickle.load(pkl_file)\n",
    "    valid_chars, maxlen, max_features = data[\"valid_chars\"], data[\"max_len\"], data[\"volcab_size\"]\n",
    "\n",
    "    # Convert characters to int and pad\n",
    "    X = [[valid_chars[y] if y in valid_chars else 0 for y in x] for x in org_X]\n",
    "    X = pad_sequences(X, maxlen=maxlen, value=0.)\n",
    "\n",
    "    # Convert labels to 0-1\n",
    "    Y = labels\n",
    "    return X, Y, maxlen, max_features\n",
    "\n",
    "\n",
    "def run():\n",
    "    testX, testY, max_len, volcab_size = get_xshell_data()\n",
    "    print(\"X len:\", len(testX), \"Y len:\", len(testY))\n",
    "    print(testX[-1:])\n",
    "    print(testY[-1:])\n",
    "    testX=np.array(testX)\n",
    "    testY=np.array(testY)\n",
    "    \n",
    "    \"\"\"\n",
    "    filename = 'dnstunnel.module'\n",
    "    loaded_model =load_model(filename)\n",
    "\n",
    "    predictions = loaded_model.predict(testX)\n",
    "    print(predictions)\n",
    "    cnt = 0\n",
    "    global org_X\n",
    "    for i,p in enumerate(predictions):\n",
    "        #if abs(p[2]-testY[i][2]) < 0.1:\n",
    "        if p[2]>p[1] and p[2]>p[0]:\n",
    "            cnt += 1\n",
    "        else:\n",
    "            print(\"found data not detected:\")\n",
    "            print(\"original subdomain:\", org_X[i])\n",
    "            print(\"prediction compare:\", p, testY[i])\n",
    "    print(\"Dectected cnt:\", cnt, \"total:\", len(predictions))\n",
    "    print(\"Dectect Rate is:\", cnt/(len(predictions)+.0))\n",
    "    \"\"\"\n",
    "    filename = 'dnstunnel.module'\n",
    "    loaded_model =load_model(filename)\n",
    "\n",
    "    predictions = loaded_model.predict(testX)\n",
    "    #print(predictions)\n",
    "    cnt = 0\n",
    "    global org_X\n",
    "    for i,p in enumerate(predictions):\n",
    "        #if abs(p[2]-testY[i][2]) < 0.1:\n",
    "        if p[2]>p[1] and p[2]>p[0]:\n",
    "            cnt += 1\n",
    "        else:\n",
    "            print(\"found data not detected:\")\n",
    "            print(\"original subdomain:\", org_X[i])\n",
    "            print(\"prediction compare:\", p, testY[i])\n",
    "    print(\"Dectected cnt:\", cnt, \"total:\", len(predictions))\n",
    "    print(\"Dectect Rate is:\", cnt/(len(predictions)+.0))\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    run()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
